package dev.langchain4j.service;

import dev.langchain4j.Internal;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.guardrail.ChatExecutor;
import dev.langchain4j.guardrail.GuardrailRequestParams;
import dev.langchain4j.guardrail.OutputGuardrailRequest;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.service.tool.ToolExecution;
import dev.langchain4j.service.tool.ToolExecutor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

import static dev.langchain4j.internal.Utils.copy;
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;

@Internal
/**
 * AI服务流式响应处理器，用于处理语言模型的流式响应。
 *
 * 主要功能包括：
 * 1. 处理模型流式输出的部分响应（通常是单个token）
 * 2. 处理工具执行请求，包括部分请求和完整请求
 * 3. 管理聊天记忆和临时记忆
 * 4. 执行输出防护（guardrails）检查
 * 5. 处理完整的聊天响应，包括工具执行和结果返回
 * 6. 错误处理和回调
 *
 * 该类实现了StreamingChatResponseHandler接口，支持处理文本响应和工具调用请求，
 * 并能与AI服务的其他组件（如工具执行器、聊天记忆、防护服务等）协调工作。
 */
class AiServiceStreamingResponseHandler implements StreamingChatResponseHandler {

    private static final Logger LOG = LoggerFactory.getLogger(AiServiceStreamingResponseHandler.class);

    /**
     * 聊天执行器，用于执行聊天请求
     */
    private final ChatExecutor chatExecutor;

    /**
     * AI服务上下文，包含模型、内存等核心组件
     */
    private final AiServiceContext context;

    /**
     * 内存ID，用于标识特定的聊天记忆
     */
    private final Object memoryId;

    /**
     * 通用防护参数，包含防护检查所需的基本信息
     */
    private final GuardrailRequestParams commonGuardrailParams;

    /**
     * 方法键，用于标识特定的AI服务方法
     */
    private final Object methodKey;

    /**
     * 部分响应处理器，处理模型流式输出的每个token
     */
    private final Consumer<String> partialResponseHandler;

    /**
     * 部分工具执行请求处理器，处理工具执行请求的部分信息
     */
    private final BiConsumer<Integer, ToolExecutionRequest> partialToolExecutionRequestHandler;

    /**
     * 完整工具执行请求处理器，处理完整的工具执行请求
     */
    private final BiConsumer<Integer, ToolExecutionRequest> completeToolExecutionRequestHandler;

    /**
     * 工具执行处理器，处理工具执行结果
     */
    private final Consumer<ToolExecution> toolExecutionHandler;

    /**
     * 完整响应处理器，处理模型的完整响应
     */
    private final Consumer<ChatResponse> completeResponseHandler;

    /**
     * 错误处理器，处理流式响应过程中的异常
     */
    private final Consumer<Throwable> errorHandler;

    /**
     * 临时内存，用于存储聊天过程中的消息
     */
    private final ChatMemory temporaryMemory;

    /**
     * 令牌使用情况，记录模型调用的令牌消耗
     */
    private final TokenUsage tokenUsage;

    /**
     * 工具规范列表，定义可用的工具及其接口
     */
    private final List<ToolSpecification> toolSpecifications;

    /**
     * 工具执行器映射，根据工具名称映射到具体的执行器实现
     */
    private final Map<String, ToolExecutor> toolExecutors;

    /**
     * 响应缓冲区，用于在输出防护检查期间暂存模型响应
     */
    private final List<String> responseBuffer = new ArrayList<>();

    /**
     * 是否启用输出防护，标识是否需要进行输出防护检查
     */
    private final boolean hasOutputGuardrails;

    AiServiceStreamingResponseHandler(
            ChatExecutor chatExecutor,
            AiServiceContext context,
            Object memoryId,
            Consumer<String> partialResponseHandler,
            BiConsumer<Integer, ToolExecutionRequest> partialToolExecutionRequestHandler,
            BiConsumer<Integer, ToolExecutionRequest> completeToolExecutionRequestHandler,
            Consumer<ToolExecution> toolExecutionHandler,
            Consumer<ChatResponse> completeResponseHandler,
            Consumer<Throwable> errorHandler,
            ChatMemory temporaryMemory,
            TokenUsage tokenUsage,
            List<ToolSpecification> toolSpecifications,
            Map<String, ToolExecutor> toolExecutors,
            GuardrailRequestParams commonGuardrailParams,
            Object methodKey) {
        this.chatExecutor = ensureNotNull(chatExecutor, "chatExecutor");
        this.context = ensureNotNull(context, "context");
        this.memoryId = ensureNotNull(memoryId, "memoryId");
        this.methodKey = methodKey;

        this.partialResponseHandler = ensureNotNull(partialResponseHandler, "partialResponseHandler");
        this.partialToolExecutionRequestHandler = partialToolExecutionRequestHandler;
        this.completeToolExecutionRequestHandler = completeToolExecutionRequestHandler;
        this.completeResponseHandler = completeResponseHandler;
        this.toolExecutionHandler = toolExecutionHandler;
        this.errorHandler = errorHandler;

        this.temporaryMemory = temporaryMemory;
        this.tokenUsage = ensureNotNull(tokenUsage, "tokenUsage");
        this.commonGuardrailParams = commonGuardrailParams;

        this.toolSpecifications = copy(toolSpecifications);
        this.toolExecutors = copy(toolExecutors);
        this.hasOutputGuardrails = context.guardrailService().hasOutputGuardrails(methodKey);
    }

    @Override
    public void onPartialResponse(String partialResponse) {
        // If we're using output guardrails, then buffer the partial response until the guardrails have completed
        if (hasOutputGuardrails) {
            responseBuffer.add(partialResponse);
        } else {
            partialResponseHandler.accept(partialResponse);
        }
    }

    @Override
    public void onPartialToolExecutionRequest(int index, ToolExecutionRequest partialToolExecutionRequest) {
        // If we're using output guardrails, then buffer the partial response until the guardrails have completed
        partialToolExecutionRequestHandler.accept(index, partialToolExecutionRequest);
    }

    /**
     * 处理模型完成的完整响应。
     * <p>
     * 主要执行流程：
     * 1. 将AI消息添加到聊天记忆中
     * 2. 检查是否存在工具执行请求：
     * - 如果存在，则依次执行所有请求的工具
     * - 将每个工具的执行结果作为新消息添加到记忆中
     * - 触发工具执行处理器（如果配置）
     * - 使用更新后的消息历史发起新的聊天请求，继续处理
     * 3. 如果没有工具请求：
     * - 执行输出防护检查（如果启用）
     * - 处理并发送缓冲的流式响应片段
     * - 调用完整的响应处理器
     *
     * @param completeResponse 模型生成的完整响应
     */
    @Override
    public void onCompleteResponse(ChatResponse completeResponse) {
        AiMessage aiMessage = completeResponse.aiMessage();
        addToMemory(aiMessage);

        if (aiMessage.hasToolExecutionRequests()) {
            // 处理工具执行请求
            for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {
                String toolName = toolExecutionRequest.name();
                ToolExecutor toolExecutor = toolExecutors.get(toolName);

                // 执行工具并获取结果
                String toolExecutionResult = toolExecutor.execute(toolExecutionRequest, memoryId);

                // 创建工具执行结果消息并添加到记忆中
                ToolExecutionResultMessage toolExecutionResultMessage =
                        ToolExecutionResultMessage.from(toolExecutionRequest, toolExecutionResult);
                addToMemory(toolExecutionResultMessage);

                // 触发工具执行处理器（如果配置）
                if (toolExecutionHandler != null) {
                    ToolExecution toolExecution = ToolExecution.builder()
                            .request(toolExecutionRequest)
                            .result(toolExecutionResult)
                            .build();
                    toolExecutionHandler.accept(toolExecution);
                }
            }

            // 使用更新后的消息历史发起新的聊天请求
            ChatRequest chatRequest = ChatRequest.builder()
                    .messages(messagesToSend(memoryId))
                    .toolSpecifications(toolSpecifications)
                    .build();

            // 创建新的响应处理器并继续处理
            var handler = new AiServiceStreamingResponseHandler(
                    chatExecutor,
                    context,
                    memoryId,
                    partialResponseHandler,
                    partialToolExecutionRequestHandler,
                    completeToolExecutionRequestHandler,
                    toolExecutionHandler,
                    completeResponseHandler,
                    errorHandler,
                    temporaryMemory,
                    TokenUsage.sum(tokenUsage, completeResponse.metadata().tokenUsage()),
                    toolSpecifications,
                    toolExecutors,
                    commonGuardrailParams,
                    methodKey);

            context.streamingChatModel.chat(chatRequest, handler);
        } else {
            // 没有工具请求，处理普通响应
            if (completeResponseHandler != null) {
                ChatResponse finalChatResponse = ChatResponse.builder()
                        .aiMessage(aiMessage)
                        .metadata(completeResponse.metadata().toBuilder()
                                .tokenUsage(tokenUsage.add(
                                        completeResponse.metadata().tokenUsage()))
                                .build())
                        .build();

                // 执行输出防护检查（如果启用）
                if (hasOutputGuardrails) {
                    if (commonGuardrailParams != null) {
                        var newCommonParams = GuardrailRequestParams.builder()
                                .chatMemory(getMemory())
                                .augmentationResult(commonGuardrailParams.augmentationResult())
                                .userMessageTemplate(commonGuardrailParams.userMessageTemplate())
                                .variables(commonGuardrailParams.variables())
                                .build();

                        var outputGuardrailParams = OutputGuardrailRequest.builder()
                                .responseFromLLM(finalChatResponse)
                                .chatExecutor(chatExecutor)
                                .requestParams(newCommonParams)
                                .build();

                        finalChatResponse =
                                context.guardrailService().executeGuardrails(methodKey, outputGuardrailParams);
                    }

                    // 处理缓冲的响应片段
                    responseBuffer.forEach(partialResponseHandler::accept);
                    responseBuffer.clear();
                }

                completeResponseHandler.accept(finalChatResponse);
            }
        }
    }

    private ChatMemory getMemory() {
        return getMemory(memoryId);
    }

    private ChatMemory getMemory(Object memId) {
        return context.hasChatMemory() ? context.chatMemoryService.getOrCreateChatMemory(memoryId) : temporaryMemory;
    }

    private void addToMemory(ChatMessage chatMessage) {
        getMemory().add(chatMessage);
    }

    private List<ChatMessage> messagesToSend(Object memoryId) {
        return getMemory(memoryId).messages();
    }

    @Override
    public void onError(Throwable error) {
        if (errorHandler != null) {
            try {
                errorHandler.accept(error);
            } catch (Exception e) {
                LOG.error("While handling the following error...", error);
                LOG.error("...the following error happened", e);
            }
        } else {
            LOG.warn("Ignored error", error);
        }
    }
}
